#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Towards and atmosphere more favourable to firestorm development in Europe """

import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from mpl_toolkits.basemap import Basemap, maskoceans
import matplotlib.colors as mcolors
import datetime as dt
from matplotlib.colors import LinearSegmentedColormap
import cmaps
from pandas import Series, DataFrame
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from matplotlib.ticker import FormatStrFormatter


font = {'size'   : 5}
matplotlib.rc('font', **font)

path_FWIKi='../2-ERA5_firestorm-risk/3-Firestorm-risk/ERA5/'
path_fuel='../3-Fuel_load/'

# Read files
FWI_Ki_days = xr.open_dataset(path_FWIKi+'FWI-Ki_ndays-present.nc').__xarray_dataarray_variable__.squeeze()  # Number of days per year with high firestorm risk
Years = xr.open_dataset(path_fuel+'Years_fuelload6.nc')['DMP']                                               # Number of years needed to produce a fuel load of 6 kg m-2
# Reduce DMP resolution
Years_interp = Years.interp(lat=Years.lat[::5], lon=Years.lon[::5], method='linear')

# Map definition
m = Basemap(projection='tmerc', resolution='h',
				llcrnrlat=30., urcrnrlat=65.,
				llcrnrlon=-10., urcrnrlon=62.,
                        lat_0=47., lon_0=15.)
lats = FWI_Ki_days['latitude']
lons = FWI_Ki_days['longitude']
lon, lat = np.meshgrid(lons.values, lats.values)
x, y = m(lon, lat)

# Plot of the number of days per year with high firestorm risk from ERA5
fig = plt.figure(figsize=(7,4))
gs = fig.add_gridspec(1,2)
ax1 = fig.add_subplot(gs[0,0])
m.drawcoastlines(linewidth=0.4,zorder=3)
m.drawcountries(linewidth=0.3,zorder=3)
m.drawparallels(np.arange(0,80,5),linewidth=0.3,labels=[1,0,0,0],zorder=3)
m.drawmeridians(np.arange(-60,80,10),linewidth=0.3,labels=[0,0,0,1],zorder=3)
m.drawlsmask(land_color='none',ocean_color='white',lakes=True,zorder=2)
clevs = np.concatenate((np.array([1.]),np.arange(5.,35.,5.)))                                       
cs2 = m.contourf(x, y, FWI_Ki_days, clevs, cmap=cmaps.WhiteBlueGreenYellowRed, alpha=0.95, extend='max',zorder=1) 
plt.title('Atmospheric firestorm risk ($days/year$)',loc='left',weight='bold',size=5,pad=4.)
ax1.text(-0.1, 1.05, 'a',family='sans-serif',weight='bold',size=7, horizontalalignment='left', verticalalignment='top',transform=ax1.transAxes)
axins = inset_axes(ax1,
                   width="3%",  
                   height="100%",  
                   loc='lower left',
                   bbox_to_anchor=(1.01,0., 1, 1),
                   bbox_transform=ax1.transAxes,
                   borderpad=0,
                   )
cbar = fig.colorbar(cs2,orientation="vertical",cax=axins,ticks=np.concatenate((np.array([1.]),np.arange(5.,35.,5.))))
plt.gcf()

# Plot of the number of years needed to produce a fuel load of 6 kg m-2
dmp_colors=np.array([[3,90,0],[20,190,20],                # We define a new colormap
                    [220,220,0], [255,153,51],
                    [255,204,153], [255,102,102],
                    [255,204,204]], np.float32)/255.0
colormap = mcolors.ListedColormap(dmp_colors)
colormap.set_over('w')
lats = Years_interp['lat']
lons = Years_interp['lon']
lon , lat = np.meshgrid(lons, lats)
x, y = m(lon, lat)
ax2 = fig.add_subplot(gs[0,1])
m.drawcoastlines(linewidth=0.4,zorder=3)
m.drawcountries(linewidth=0.3,zorder=3)
m.drawparallels(np.arange(0,80,5),linewidth=0.3,labels=[1,0,0,0],zorder=3)
m.drawmeridians(np.arange(-60,80,10),linewidth=0.3,labels=[0,0,0,1],zorder=3)
m.drawlsmask(land_color='none',ocean_color='white',lakes=True,zorder=2)
clevs = np.array([0.,2.,4.,6.,8.,10.,15.,25.])  
norm = mcolors.BoundaryNorm(clevs, colormap.N)
cs2 = m.contourf(x, y, Years_interp, clevs, cmap=colormap, alpha=0.95, norm=norm, extend='max')
plt.title('Fuel production time ($years$)',loc='left',weight='bold',size=5,pad=4.)
ax2.text(-0.1, 1.05, 'b',family='sans-serif',weight='bold',size=7, horizontalalignment='left', verticalalignment='top',transform=ax2.transAxes)
axins = inset_axes(ax2,
                   width="3%",  
                   height="100%",  
                   loc='lower left',
                   bbox_to_anchor=(1.01,0., 1, 1),
                   bbox_transform=ax2.transAxes,
                   borderpad=0,
                   )
cbar = fig.colorbar(cs2,orientation="vertical",cax=axins,ticks=clevs)
plt.gcf()

# Save figure
plt.subplots_adjust(top = 0.9, bottom=0.1, hspace=0., wspace=0.25)
fig.savefig('Figure3.png',dpi=300,bbox_inches='tight')   
fig.savefig('Figure3.pdf',bbox_inches='tight')   
